{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Enhancing RAG with Contextual Retrieval\n",
    "\n",
    "> Note: For more background information on Contextual Retrieval, including additional performance evaluations on various datasets, we recommend reading our accompanying  [blog post](https://www.anthropic.com/news/contextual-retrieval).\n",
    "\n",
    "Retrieval Augmented Generation (RAG) enables Claude to leverage your internal knowledge bases, codebases, or any other corpus of documents when providing a response. Enterprises are increasingly building RAG applications to improve workflows in customer support, Q&A over internal company documents, financial & legal analysis, code generation, and much more.\n",
    "\n",
    "In a [separate guide](https://github.com/anthropics/anthropic-cookbook/blob/main/capabilities/retrieval_augmented_generation/guide.ipynb), we walked through setting up a basic retrieval system, demonstrated how to evaluate its performance, and then outlined a few techniques to improve performance. In this guide, we present a technique for improving retrieval performance: Contextual Embeddings.\n",
    "\n",
    "In traditional RAG, documents are typically split into smaller chunks for efficient retrieval. While this approach works well for many applications, it can lead to problems when individual chunks lack sufficient context. Contextual Embeddings solve this problem by adding relevant context to each chunk before embedding. This method improves the quality of each embedded chunk, allowing for more accurate retrieval and thus better overall performance. Averaged across all data sources we tested, Contextual Embeddings reduced the top-20-chunk retrieval failure rate by 35%.\n",
    "\n",
    "The same chunk-specific context can also be used with BM25 search to further improve retrieval performance. We introduce this technique in the \"Contextual BM25\" section.\n",
    "\n",
    "In this guide, we'll demonstrate how to build and optimize a Contextual Retrieval system using a dataset of 9 codebases as our knowledge base. We'll walk through:\n",
    "\n",
    "1) Setting up a basic retrieval pipeline to establish a baseline for performance.\n",
    "\n",
    "2) Contextual Embeddings: what it is, why it works, and how prompt caching makes it practical for production use cases.\n",
    "\n",
    "3) Implementing Contextual Embeddings and demonstrating performance improvements.\n",
    "\n",
    "4) Contextual BM25: improving performance with *contextual* BM25 hybrid search.\n",
    "\n",
    "5) Improving performance with reranking,\n",
    "\n",
    "### Evaluation Metrics & Dataset:\n",
    "\n",
    "We use a pre-chunked dataset of 9 codebases - all of which have been chunked according to a basic character splitting mechanism. Our evaluation dataset contains 248 queries - each of which contains a 'golden chunk.' We'll use a metric called Pass@k to evaluate performance. Pass@k checks whether or not the 'golden document' was present in the first k documents retrieved for each query. Contextual Embeddings in this case helped us to improve Pass@10 performance from ~87% --> ~95%.\n",
    "\n",
    "You can find the code files and their chunks in `data/codebase_chunks.json` and the evaluation dataset in `data/evaluation_set.jsonl`\n",
    "\n",
    "#### Additional Notes:\n",
    "\n",
    "Prompt caching is helpful in managing costs when using this retrieval method. This feature is currently available on Anthropic's first-party API, and is coming soon to our third-party partner environments in AWS Bedrock and GCP Vertex. We know that many of our customers leverage AWS Knowledge Bases and GCP Vertex AI APIs when building RAG solutions, and this method can be used on either platform with a bit of customization. Consider reaching out to Anthropic or your AWS/GCP account team for guidance on this!\n",
    "\n",
    "To make it easier to use this method on Bedrock, the AWS team has provided us with code that you can use to implement a Lambda function that adds context to each document. If you deploy this Lambda function, you can select it as a custom chunking option when configuring a [Bedrock Knowledge Base](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-create.html). You can find this code in `contextual-rag-lambda-function`. The main lambda function code is in `lambda_function.py`.\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "1) Setup\n",
    "\n",
    "2) Basic RAG\n",
    "\n",
    "3) Contextual Embeddings\n",
    "\n",
    "4) Contextual BM25\n",
    "\n",
    "5) Reranking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Before starting this guide, ensure you have:\n",
    "\n",
    "**Technical Skills:**\n",
    "- Intermediate Python programming\n",
    "- Basic understanding of RAG (Retrieval Augmented Generation)\n",
    "- Familiarity with vector databases and embeddings\n",
    "- Basic command-line proficiency\n",
    "\n",
    "**System Requirements:**\n",
    "- Python 3.8+\n",
    "- Docker installed and running (optional, for BM25 search)\n",
    "- 4GB+ available RAM\n",
    "- ~5-10 GB disk space for vector databases\n",
    "\n",
    "**API Access:**\n",
    "- [Anthropic API key](https://console.anthropic.com/) (free tier sufficient)\n",
    "- [Voyage AI API key](https://www.voyageai.com/)\n",
    "- [Cohere API key](https://cohere.com/)\n",
    "\n",
    "**Time & Cost:**\n",
    "- Expected completion time: 30-45 minutes\n",
    "- API costs: ~$5-10 to run through the full dataset\n",
    "\n",
    "### Libraries \n",
    "\n",
    "We'll need a few libraries, including:\n",
    "\n",
    "1) `anthropic` - to interact with Claude\n",
    "\n",
    "2) `voyageai` - to generate high quality embeddings\n",
    "\n",
    "3) `cohere` - for reranking\n",
    "\n",
    "4) `elasticsearch` for performant BM25 search\n",
    "\n",
    "3) `pandas`, `numpy`, `matplotlib`, and `scikit-learn` for data manipulation and visualization\n",
    "\n",
    "### Environment Variables \n",
    "\n",
    "Ensure the following environment variables are set:\n",
    "\n",
    "```\n",
    "- VOYAGE_API_KEY\n",
    "- ANTHROPIC_API_KEY\n",
    "- COHERE_API_KEY\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install --upgrade anthropic voyageai cohere elasticsearch pandas numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define our model names up front to make it easier to change models as new models are released"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"claude-haiku-4-5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by initializing the Anthropic client that we'll use for generating contextual descriptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import anthropic\n",
    "\n",
    "client = anthropic.Anthropic(\n",
    "    # This is the default and can be omitted\n",
    "    api_key=os.getenv(\"ANTHROPIC_API_KEY\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize a Vector DB Class\n",
    "\n",
    "We'll create a VectorDB class to handle embedding storage and similarity search. This class serves three key functions in our RAG pipeline:\n",
    "\n",
    "1. **Embedding Generation**: Converts text chunks into vector representations using Voyage AI's embedding model\n",
    "2. **Storage & Caching**: Saves embeddings to disk to avoid re-computing them (which saves time and API costs)\n",
    "3. **Similarity Search**: Retrieves the most relevant chunks for a given query using cosine similarity\n",
    "\n",
    "\n",
    "For this guide, we're using a simple in-memory vector database with pickle serialization. This makes the code easy to understand and requires no external dependencies. The class automatically saves embeddings to disk after generation, so you only pay the embedding cost once. \n",
    "\n",
    "For production use, consider hosted vector database solutions.\n",
    "\n",
    "The VectorDB class below follows the same interface patterns you'd use with production solutions, making it easy to swap out later. Key features include batch processing (128 chunks at a time), progress tracking with tqdm, and query caching to speed up repeated searches during evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pickle\n",
    "from typing import Any\n",
    "\n",
    "import numpy as np\n",
    "import voyageai\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "class VectorDB:\n",
    "    def __init__(self, name: str, api_key=None):\n",
    "        if api_key is None:\n",
    "            api_key = os.getenv(\"VOYAGE_API_KEY\")\n",
    "        self.client = voyageai.Client(api_key=api_key)\n",
    "        self.name = name\n",
    "        self.embeddings = []\n",
    "        self.metadata = []\n",
    "        self.query_cache = {}\n",
    "        self.db_path = f\"./data/{name}/vector_db.pkl\"\n",
    "\n",
    "    def load_data(self, dataset: list[dict[str, Any]]):\n",
    "        if self.embeddings and self.metadata:\n",
    "            print(\"Vector database is already loaded. Skipping data loading.\")\n",
    "            return\n",
    "        if os.path.exists(self.db_path):\n",
    "            print(\"Loading vector database from disk.\")\n",
    "            self.load_db()\n",
    "            return\n",
    "\n",
    "        texts_to_embed = []\n",
    "        metadata = []\n",
    "        total_chunks = sum(len(doc[\"chunks\"]) for doc in dataset)\n",
    "\n",
    "        with tqdm(total=total_chunks, desc=\"Processing chunks\") as pbar:\n",
    "            for doc in dataset:\n",
    "                for chunk in doc[\"chunks\"]:\n",
    "                    texts_to_embed.append(chunk[\"content\"])\n",
    "                    metadata.append(\n",
    "                        {\n",
    "                            \"doc_id\": doc[\"doc_id\"],\n",
    "                            \"original_uuid\": doc[\"original_uuid\"],\n",
    "                            \"chunk_id\": chunk[\"chunk_id\"],\n",
    "                            \"original_index\": chunk[\"original_index\"],\n",
    "                            \"content\": chunk[\"content\"],\n",
    "                        }\n",
    "                    )\n",
    "                    pbar.update(1)\n",
    "\n",
    "        self._embed_and_store(texts_to_embed, metadata)\n",
    "        self.save_db()\n",
    "\n",
    "        print(f\"Vector database loaded and saved. Total chunks processed: {len(texts_to_embed)}\")\n",
    "\n",
    "    def _embed_and_store(self, texts: list[str], data: list[dict[str, Any]]):\n",
    "        batch_size = 128\n",
    "        with tqdm(total=len(texts), desc=\"Embedding chunks\") as pbar:\n",
    "            result = []\n",
    "            for i in range(0, len(texts), batch_size):\n",
    "                batch = texts[i : i + batch_size]\n",
    "                batch_result = self.client.embed(batch, model=\"voyage-2\").embeddings\n",
    "                result.extend(batch_result)\n",
    "                pbar.update(len(batch))\n",
    "\n",
    "        self.embeddings = result\n",
    "        self.metadata = data\n",
    "\n",
    "    def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:\n",
    "        if query in self.query_cache:\n",
    "            query_embedding = self.query_cache[query]\n",
    "        else:\n",
    "            query_embedding = self.client.embed([query], model=\"voyage-2\").embeddings[0]\n",
    "            self.query_cache[query] = query_embedding\n",
    "\n",
    "        if not self.embeddings:\n",
    "            raise ValueError(\"No data loaded in the vector database.\")\n",
    "\n",
    "        similarities = np.dot(self.embeddings, query_embedding)\n",
    "        top_indices = np.argsort(similarities)[::-1][:k]\n",
    "\n",
    "        top_results = []\n",
    "        for idx in top_indices:\n",
    "            result = {\n",
    "                \"metadata\": self.metadata[idx],\n",
    "                \"similarity\": float(similarities[idx]),\n",
    "            }\n",
    "            top_results.append(result)\n",
    "\n",
    "        return top_results\n",
    "\n",
    "    def save_db(self):\n",
    "        data = {\n",
    "            \"embeddings\": self.embeddings,\n",
    "            \"metadata\": self.metadata,\n",
    "            \"query_cache\": json.dumps(self.query_cache),\n",
    "        }\n",
    "        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)\n",
    "        with open(self.db_path, \"wb\") as file:\n",
    "            pickle.dump(data, file)\n",
    "\n",
    "    def load_db(self):\n",
    "        if not os.path.exists(self.db_path):\n",
    "            raise ValueError(\n",
    "                \"Vector database file not found. Use load_data to create a new database.\"\n",
    "            )\n",
    "        with open(self.db_path, \"rb\") as file:\n",
    "            data = pickle.load(file)\n",
    "        self.embeddings = data[\"embeddings\"]\n",
    "        self.metadata = data[\"metadata\"]\n",
    "        self.query_cache = json.loads(data[\"query_cache\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use this class to load our dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing chunks: 100%|██████████| 737/737 [00:00<00:00, 985400.72it/s]\n",
      "Embedding chunks: 100%|██████████| 737/737 [00:42<00:00, 17.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vector database loaded and saved. Total chunks processed: 737\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Load your transformed dataset\n",
    "with open(\"data/codebase_chunks.json\") as f:\n",
    "    transformed_dataset = json.load(f)\n",
    "\n",
    "# Initialize the VectorDB\n",
    "base_db = VectorDB(\"base_db\")\n",
    "\n",
    "# Load and process the data\n",
    "base_db.load_data(transformed_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic RAG\n",
    "\n",
    "To get started, we'll set up a basic RAG pipeline using a bare bones approach. This is sometimes called 'Naive RAG' by many in the industry. A basic RAG pipeline includes the following 3 steps:\n",
    "\n",
    "1) Chunk documents by heading - containing only the content from each subheading\n",
    "\n",
    "2) Embed each document\n",
    "\n",
    "3) Use Cosine similarity to retrieve documents in order to answer query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from collections.abc import Callable\n",
    "from typing import Any\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def load_jsonl(file_path: str) -> list[dict[str, Any]]:\n",
    "    \"\"\"Load JSONL file and return a list of dictionaries.\"\"\"\n",
    "    with open(file_path) as file:\n",
    "        return [json.loads(line) for line in file]\n",
    "\n",
    "\n",
    "def evaluate_retrieval(\n",
    "    queries: list[dict[str, Any]], retrieval_function: Callable, db, k: int = 20\n",
    ") -> dict[str, float]:\n",
    "    total_score = 0\n",
    "    total_queries = len(queries)\n",
    "\n",
    "    for query_item in tqdm(queries, desc=\"Evaluating retrieval\"):\n",
    "        query = query_item[\"query\"]\n",
    "        golden_chunk_uuids = query_item[\"golden_chunk_uuids\"]\n",
    "\n",
    "        # Find all golden chunk contents\n",
    "        golden_contents = []\n",
    "        for doc_uuid, chunk_index in golden_chunk_uuids:\n",
    "            golden_doc = next(\n",
    "                (doc for doc in query_item[\"golden_documents\"] if doc[\"uuid\"] == doc_uuid), None\n",
    "            )\n",
    "            if not golden_doc:\n",
    "                print(f\"Warning: Golden document not found for UUID {doc_uuid}\")\n",
    "                continue\n",
    "\n",
    "            golden_chunk = next(\n",
    "                (chunk for chunk in golden_doc[\"chunks\"] if chunk[\"index\"] == chunk_index), None\n",
    "            )\n",
    "            if not golden_chunk:\n",
    "                print(\n",
    "                    f\"Warning: Golden chunk not found for index {chunk_index} in document {doc_uuid}\"\n",
    "                )\n",
    "                continue\n",
    "\n",
    "            golden_contents.append(golden_chunk[\"content\"].strip())\n",
    "\n",
    "        if not golden_contents:\n",
    "            print(f\"Warning: No golden contents found for query: {query}\")\n",
    "            continue\n",
    "\n",
    "        retrieved_docs = retrieval_function(query, db, k=k)\n",
    "\n",
    "        # Count how many golden chunks are in the top k retrieved documents\n",
    "        chunks_found = 0\n",
    "        for golden_content in golden_contents:\n",
    "            for doc in retrieved_docs[:k]:\n",
    "                retrieved_content = (\n",
    "                    doc[\"metadata\"]\n",
    "                    .get(\"original_content\", doc[\"metadata\"].get(\"content\", \"\"))\n",
    "                    .strip()\n",
    "                )\n",
    "                if retrieved_content == golden_content:\n",
    "                    chunks_found += 1\n",
    "                    break\n",
    "\n",
    "        query_score = chunks_found / len(golden_contents)\n",
    "        total_score += query_score\n",
    "\n",
    "    average_score = total_score / total_queries\n",
    "    pass_at_n = average_score * 100\n",
    "    return {\"pass_at_n\": pass_at_n, \"average_score\": average_score, \"total_queries\": total_queries}\n",
    "\n",
    "\n",
    "def retrieve_base(query: str, db, k: int = 20) -> list[dict[str, Any]]:\n",
    "    \"\"\"\n",
    "    Retrieve relevant documents using either VectorDB or ContextualVectorDB.\n",
    "\n",
    "    :param query: The query string\n",
    "    :param db: The VectorDB or ContextualVectorDB instance\n",
    "    :param k: Number of top results to retrieve\n",
    "    :return: List of retrieved documents\n",
    "    \"\"\"\n",
    "    return db.search(query, k=k)\n",
    "\n",
    "\n",
    "def evaluate_db(db, original_jsonl_path: str, k):\n",
    "    # Load the original JSONL data for queries and ground truth\n",
    "    original_data = load_jsonl(original_jsonl_path)\n",
    "\n",
    "    # Evaluate retrieval\n",
    "    results = evaluate_retrieval(original_data, retrieve_base, db, k)\n",
    "    return results\n",
    "\n",
    "\n",
    "def evaluate_and_display(db, jsonl_path: str, k_values: list[int] = None, db_name: str = \"\"):\n",
    "    \"\"\"\n",
    "    Evaluate retrieval performance across multiple k values and display formatted results.\n",
    "\n",
    "    Args:\n",
    "        db: Vector database instance (VectorDB or ContextualVectorDB)\n",
    "        jsonl_path: Path to evaluation dataset\n",
    "        k_values: List of k values to evaluate (default: [5, 10, 20])\n",
    "        db_name: Optional name for the database being evaluated\n",
    "\n",
    "    Returns:\n",
    "        Dict mapping k values to their results\n",
    "    \"\"\"\n",
    "    if k_values is None:\n",
    "        k_values = [5, 10, 20]\n",
    "    results = {}\n",
    "\n",
    "    print(f\"{'=' * 60}\")\n",
    "    if db_name:\n",
    "        print(f\"Evaluation Results: {db_name}\")\n",
    "    else:\n",
    "        print(\"Evaluation Results\")\n",
    "    print(f\"{'=' * 60}\\n\")\n",
    "\n",
    "    for k in k_values:\n",
    "        print(f\"Evaluating Pass@{k}...\")\n",
    "        results[k] = evaluate_db(db, jsonl_path, k)\n",
    "        print()  # Add spacing between evaluations\n",
    "\n",
    "    # Print summary table\n",
    "    print(f\"{'=' * 60}\")\n",
    "    print(f\"{'Metric':<15} {'Pass Rate':<15} {'Score':<15}\")\n",
    "    print(f\"{'-' * 60}\")\n",
    "    for k in k_values:\n",
    "        pass_rate = f\"{results[k]['pass_at_n']:.2f}%\"\n",
    "        score = f\"{results[k]['average_score']:.4f}\"\n",
    "        print(f\"{'Pass@' + str(k):<15} {pass_rate:<15} {score:<15}\")\n",
    "    print(f\"{'=' * 60}\\n\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's establish our baseline performance by evaluating the basic RAG system. We'll test at k=5, 10, and 20 to see how many of the golden chunks appear in the top retrieved results. This gives us a benchmark to measure improvement against.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================================================\n",
      "Evaluation Results: Contextual Embeddings\n",
      "============================================================\n",
      "\n",
      "Evaluating Pass@5...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 65.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluating Pass@10...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluating Pass@20...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============================================================\n",
      "Metric          Pass Rate       Score          \n",
      "------------------------------------------------------------\n",
      "Pass@5          80.92%          0.8092         \n",
      "Pass@10         87.15%          0.8715         \n",
      "Pass@20         90.06%          0.9006         \n",
      "============================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results = evaluate_and_display(\n",
    "    base_db, \"data/evaluation_set.jsonl\", k_values=[5, 10, 20], db_name=\"Baseline RAG\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These results show our baseline RAG performance. The system successfully retrieves the correct chunk 81% of the time in the top 5 results, improving to 87% in the top 10, and 90% in the top 20."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contextual Embeddings\n",
    "\n",
    "With basic RAG, individual chunks often lack sufficient context when embedded in isolation. Contextual Embeddings solve this by using Claude to generate a brief description that \"situates\" each chunk within its source document. We then embed the chunk together with this context, creating richer vector representations. \n",
    "\n",
    "For each chunk in our codebase dataset, we pass both the chunk and its full source file to Claude. Claude generates a concise explanation of what the chunk contains and where it fits in the overall file. This context gets prepended to the chunk before embedding.\n",
    "\n",
    "\n",
    "### Cost and Latency Considerations\n",
    "\n",
    "**When does this cost occur?** The contextualization happens once at ingestion time, not during every query. Unlike techniques like HyDE (hypothetical document embeddings) that add latency to each search, contextual embeddings are a one-time cost when building your vector database. Prompt caching makes this practical. Since we process all chunks from the same document sequentially, we can leverage prompt caching for significant savings.\n",
    "\n",
    "1. First chunk: We write the full document to cache (pay a small premium)\n",
    "2. Subsequent chunks: Read the document from cache (90% discount on those tokens)\n",
    "3. Cache lasts 5 minutes, plenty of time to process all chunks in a document\n",
    "\n",
    "**Cost example**: For 800-token chunks in 8k-token documents with 100 tokens of generated context, the total cost is $1.02 per million document tokens. You'll see the cache savings in the logs when you run the code below.\n",
    "\n",
    "**Note:** Some embedding models have fixed input token limits. If you see worse performance with contextual embeddings, your contextualized chunks may be getting truncated—consider using an embedding model with a larger context window."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--- \n",
    "\n",
    "Let's see an example of how contextual embeddings work by generating context for a single chunk. We'll use Claude to create a situating context, and you'll also see the prompt caching metrics in action.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Situated context: This chunk contains the module documentation and initial struct definition for a differential fuzzing executor. It introduces the `DiffExecutor` struct that wraps two executors (primary and secondary) to run them sequentially with the same input, comparing their behavior for differential testing. The chunk establishes the core data structure and imports needed for the differential fuzzing implementation.\n",
      "----------\n",
      "Input tokens: 3412\n",
      "Output tokens: 76\n",
      "Cache creation input tokens: 0\n",
      "Cache read input tokens: 0\n"
     ]
    }
   ],
   "source": [
    "DOCUMENT_CONTEXT_PROMPT = \"\"\"\n",
    "<document>\n",
    "{doc_content}\n",
    "</document>\n",
    "\"\"\"\n",
    "\n",
    "CHUNK_CONTEXT_PROMPT = \"\"\"\n",
    "Here is the chunk we want to situate within the whole document\n",
    "<chunk>\n",
    "{chunk_content}\n",
    "</chunk>\n",
    "\n",
    "Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.\n",
    "Answer only with the succinct context and nothing else.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def situate_context(doc: str, chunk: str) -> str:\n",
    "    response = client.messages.create(\n",
    "        model=MODEL_NAME,\n",
    "        max_tokens=1024,\n",
    "        temperature=0.0,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc),\n",
    "                        \"cache_control\": {\n",
    "                            \"type\": \"ephemeral\"\n",
    "                        },  # we will make use of prompt caching for the full documents\n",
    "                    },\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk),\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ],\n",
    "    )\n",
    "    return response\n",
    "\n",
    "\n",
    "jsonl_data = load_jsonl(\"data/evaluation_set.jsonl\")\n",
    "# Example usage\n",
    "doc_content = jsonl_data[0][\"golden_documents\"][0][\"content\"]\n",
    "chunk_content = jsonl_data[0][\"golden_chunks\"][0][\"content\"]\n",
    "\n",
    "response = situate_context(doc_content, chunk_content)\n",
    "print(f\"Situated context: {response.content[0].text}\")\n",
    "print(\"-\" * 10)\n",
    "# Print cache performance metrics\n",
    "print(f\"Input tokens: {response.usage.input_tokens}\")\n",
    "print(f\"Output tokens: {response.usage.output_tokens}\")\n",
    "print(f\"Cache creation input tokens: {response.usage.cache_creation_input_tokens}\")\n",
    "print(f\"Cache read input tokens: {response.usage.cache_read_input_tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building the Contextual Vector Database\n",
    "\n",
    "Now that we've seen how to generate contextual descriptions for individual chunks, let's scale this up to process our entire dataset. The `ContextualVectorDB` class below extends our basic `VectorDB` with automatic contextualization during ingestion.\n",
    "\n",
    "**Key features:**\n",
    "\n",
    "- **Parallel processing**: Uses ThreadPoolExecutor to contextualize multiple chunks simultaneously (configurable thread count)\n",
    "- **Automatic prompt caching**: Processes chunks document-by-document to maximize cache hits\n",
    "- **Token tracking**: Monitors cache performance and calculates actual cost savings\n",
    "- **Persistent storage**: Saves both embeddings and contextualized metadata to disk\n",
    "\n",
    "When you run this, pay attention to the token usage statistics—you'll see that 70-80% of input tokens are read from cache, demonstrating the dramatic cost savings from prompt caching. On our 737-chunk dataset, this reduces what would be a ~$15 ingestion job down to ~$3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import threading\n",
    "import time\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from typing import Any\n",
    "\n",
    "import anthropic\n",
    "import numpy as np\n",
    "import voyageai\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "class ContextualVectorDB:\n",
    "    def __init__(self, name: str, voyage_api_key=None, anthropic_api_key=None):\n",
    "        if voyage_api_key is None:\n",
    "            voyage_api_key = os.getenv(\"VOYAGE_API_KEY\")\n",
    "        if anthropic_api_key is None:\n",
    "            anthropic_api_key = os.getenv(\"ANTHROPIC_API_KEY\")\n",
    "\n",
    "        self.voyage_client = voyageai.Client(api_key=voyage_api_key)\n",
    "        self.anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)\n",
    "        self.name = name\n",
    "        self.embeddings = []\n",
    "        self.metadata = []\n",
    "        self.query_cache = {}\n",
    "        self.db_path = f\"./data/{name}/contextual_vector_db.pkl\"\n",
    "\n",
    "        self.token_counts = {\"input\": 0, \"output\": 0, \"cache_read\": 0, \"cache_creation\": 0}\n",
    "        self.token_lock = threading.Lock()\n",
    "\n",
    "    def situate_context(self, doc: str, chunk: str) -> tuple[str, Any]:\n",
    "        DOCUMENT_CONTEXT_PROMPT = \"\"\"\n",
    "        <document>\n",
    "        {doc_content}\n",
    "        </document>\n",
    "        \"\"\"\n",
    "\n",
    "        CHUNK_CONTEXT_PROMPT = \"\"\"\n",
    "        Here is the chunk we want to situate within the whole document\n",
    "        <chunk>\n",
    "        {chunk_content}\n",
    "        </chunk>\n",
    "\n",
    "        Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.\n",
    "        Answer only with the succinct context and nothing else.\n",
    "        \"\"\"\n",
    "\n",
    "        response = self.anthropic_client.messages.create(\n",
    "            model=MODEL_NAME,\n",
    "            max_tokens=1000,\n",
    "            temperature=0.0,\n",
    "            messages=[\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": [\n",
    "                        {\n",
    "                            \"type\": \"text\",\n",
    "                            \"text\": DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc),\n",
    "                            \"cache_control\": {\n",
    "                                \"type\": \"ephemeral\"\n",
    "                            },  # we will make use of prompt caching for the full documents\n",
    "                        },\n",
    "                        {\n",
    "                            \"type\": \"text\",\n",
    "                            \"text\": CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk),\n",
    "                        },\n",
    "                    ],\n",
    "                },\n",
    "            ],\n",
    "            extra_headers={\"anthropic-beta\": \"prompt-caching-2024-07-31\"},\n",
    "        )\n",
    "        return response.content[0].text, response.usage\n",
    "\n",
    "    def load_data(self, dataset: list[dict[str, Any]], parallel_threads: int = 1):\n",
    "        if self.embeddings and self.metadata:\n",
    "            print(\"Vector database is already loaded. Skipping data loading.\")\n",
    "            return\n",
    "        if os.path.exists(self.db_path):\n",
    "            print(\"Loading vector database from disk.\")\n",
    "            self.load_db()\n",
    "            return\n",
    "\n",
    "        texts_to_embed = []\n",
    "        metadata = []\n",
    "        total_chunks = sum(len(doc[\"chunks\"]) for doc in dataset)\n",
    "\n",
    "        def process_chunk(doc, chunk):\n",
    "            # for each chunk, produce the context\n",
    "            contextualized_text, usage = self.situate_context(doc[\"content\"], chunk[\"content\"])\n",
    "            with self.token_lock:\n",
    "                self.token_counts[\"input\"] += usage.input_tokens\n",
    "                self.token_counts[\"output\"] += usage.output_tokens\n",
    "                self.token_counts[\"cache_read\"] += usage.cache_read_input_tokens\n",
    "                self.token_counts[\"cache_creation\"] += usage.cache_creation_input_tokens\n",
    "\n",
    "            return {\n",
    "                # append the context to the original text chunk\n",
    "                \"text_to_embed\": f\"{chunk['content']}\\n\\n{contextualized_text}\",\n",
    "                \"metadata\": {\n",
    "                    \"doc_id\": doc[\"doc_id\"],\n",
    "                    \"original_uuid\": doc[\"original_uuid\"],\n",
    "                    \"chunk_id\": chunk[\"chunk_id\"],\n",
    "                    \"original_index\": chunk[\"original_index\"],\n",
    "                    \"original_content\": chunk[\"content\"],\n",
    "                    \"contextualized_content\": contextualized_text,\n",
    "                },\n",
    "            }\n",
    "\n",
    "        print(f\"Processing {total_chunks} chunks with {parallel_threads} threads\")\n",
    "        with ThreadPoolExecutor(max_workers=parallel_threads) as executor:\n",
    "            futures = []\n",
    "            for doc in dataset:\n",
    "                for chunk in doc[\"chunks\"]:\n",
    "                    futures.append(executor.submit(process_chunk, doc, chunk))\n",
    "\n",
    "            for future in tqdm(as_completed(futures), total=total_chunks, desc=\"Processing chunks\"):\n",
    "                result = future.result()\n",
    "                texts_to_embed.append(result[\"text_to_embed\"])\n",
    "                metadata.append(result[\"metadata\"])\n",
    "\n",
    "        self._embed_and_store(texts_to_embed, metadata)\n",
    "        self.save_db()\n",
    "\n",
    "        # logging token usage\n",
    "        print(\n",
    "            f\"Contextual Vector database loaded and saved. Total chunks processed: {len(texts_to_embed)}\"\n",
    "        )\n",
    "        print(f\"Total input tokens without caching: {self.token_counts['input']}\")\n",
    "        print(f\"Total output tokens: {self.token_counts['output']}\")\n",
    "        print(f\"Total input tokens written to cache: {self.token_counts['cache_creation']}\")\n",
    "        print(f\"Total input tokens read from cache: {self.token_counts['cache_read']}\")\n",
    "\n",
    "        total_tokens = (\n",
    "            self.token_counts[\"input\"]\n",
    "            + self.token_counts[\"cache_read\"]\n",
    "            + self.token_counts[\"cache_creation\"]\n",
    "        )\n",
    "        savings_percentage = (\n",
    "            (self.token_counts[\"cache_read\"] / total_tokens) * 100 if total_tokens > 0 else 0\n",
    "        )\n",
    "        print(\n",
    "            f\"Total input token savings from prompt caching: {savings_percentage:.2f}% of all input tokens used were read from cache.\"\n",
    "        )\n",
    "        print(\"Tokens read from cache come at a 90 percent discount!\")\n",
    "\n",
    "    # we use voyage AI here for embeddings. Read more here: https://docs.voyageai.com/docs/embeddings\n",
    "    def _embed_and_store(self, texts: list[str], data: list[dict[str, Any]]):\n",
    "        batch_size = 128\n",
    "        result = [\n",
    "            self.voyage_client.embed(texts[i : i + batch_size], model=\"voyage-2\").embeddings\n",
    "            for i in range(0, len(texts), batch_size)\n",
    "        ]\n",
    "        self.embeddings = [embedding for batch in result for embedding in batch]\n",
    "        self.metadata = data\n",
    "\n",
    "    def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:\n",
    "        if query in self.query_cache:\n",
    "            query_embedding = self.query_cache[query]\n",
    "        else:\n",
    "            query_embedding = self.voyage_client.embed([query], model=\"voyage-2\").embeddings[0]\n",
    "            self.query_cache[query] = query_embedding\n",
    "\n",
    "        if not self.embeddings:\n",
    "            raise ValueError(\"No data loaded in the vector database.\")\n",
    "\n",
    "        similarities = np.dot(self.embeddings, query_embedding)\n",
    "        top_indices = np.argsort(similarities)[::-1][:k]\n",
    "\n",
    "        top_results = []\n",
    "        for idx in top_indices:\n",
    "            result = {\n",
    "                \"metadata\": self.metadata[idx],\n",
    "                \"similarity\": float(similarities[idx]),\n",
    "            }\n",
    "            top_results.append(result)\n",
    "        return top_results\n",
    "\n",
    "    def save_db(self):\n",
    "        data = {\n",
    "            \"embeddings\": self.embeddings,\n",
    "            \"metadata\": self.metadata,\n",
    "            \"query_cache\": json.dumps(self.query_cache),\n",
    "        }\n",
    "        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)\n",
    "        with open(self.db_path, \"wb\") as file:\n",
    "            pickle.dump(data, file)\n",
    "\n",
    "    def load_db(self):\n",
    "        if not os.path.exists(self.db_path):\n",
    "            raise ValueError(\n",
    "                \"Vector database file not found. Use load_data to create a new database.\"\n",
    "            )\n",
    "        with open(self.db_path, \"rb\") as file:\n",
    "            data = pickle.load(file)\n",
    "        self.embeddings = data[\"embeddings\"]\n",
    "        self.metadata = data[\"metadata\"]\n",
    "        self.query_cache = json.loads(data[\"query_cache\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 737 chunks with 5 threads\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing chunks: 100%|██████████| 737/737 [05:32<00:00,  2.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Contextual Vector database loaded and saved. Total chunks processed: 737\n",
      "Total input tokens without caching: 1223730\n",
      "Total output tokens: 58161\n",
      "Total input tokens written to cache: 176079\n",
      "Total input tokens read from cache: 2267069\n",
      "Total input token savings from prompt caching: 61.83% of all input tokens used were read from cache.\n",
      "Tokens read from cache come at a 90 percent discount!\n"
     ]
    }
   ],
   "source": [
    "# Load the transformed dataset\n",
    "with open(\"data/codebase_chunks.json\") as f:\n",
    "    transformed_dataset = json.load(f)\n",
    "\n",
    "# Initialize the ContextualVectorDB\n",
    "contextual_db = ContextualVectorDB(\"my_contextual_db\")\n",
    "\n",
    "# Load and process the data\n",
    "# note: consider increasing the number of parallel threads to run this faster, or reducing the number of parallel threads if concerned about hitting your API rate limit\n",
    "contextual_db.load_data(transformed_dataset, parallel_threads=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These numbers reveal the power of prompt caching for contextual embeddings:\n",
    "\n",
    "- We processed **737 chunks** across 9 codebase files\n",
    "- **61.83% of input tokens** were read from cache (2.27M tokens at 90% discount)\n",
    "- Without caching, this would cost **~$9.20** in input tokens\n",
    "- With caching, the actual cost drops to **~$2.85** (69% savings)\n",
    "\n",
    "The cache hit rate depends on how many chunks each document contains. Files with more chunks benefit more from caching since we write the full document to cache once, then read it repeatedly for each chunk in that file. This is why processing documents sequentially (rather than randomly shuffling chunks) is crucial for maximizing cache efficiency.\n",
    "\n",
    "Now let's evaluate how much this contextualization improves our retrieval performance compared to the baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================================================\n",
      "Evaluation Results: Contextual Embeddings\n",
      "============================================================\n",
      "\n",
      "Evaluating Pass@5...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluating Pass@10...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluating Pass@20...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating retrieval: 100%|██████████| 248/248 [00:03<00:00, 64.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============================================================\n",
      "Metric          Pass Rate       Score          \n",
      "------------------------------------------------------------\n",
      "Pass@5          88.12%          0.8812         \n",
      "Pass@10         92.34%          0.9234         \n",
      "Pass@20         94.29%          0.9429         \n",
      "============================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results = evaluate_and_display(\n",
    "    contextual_db,\n",
    "    \"data/evaluation_set.jsonl\",\n",
    "    k_values=[5, 10, 20],\n",
    "    db_name=\"Contextual Embeddings\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By adding context to each chunk before embedding, we've reduced retrieval failures by **~30-40%** across all k values. This means fewer irrelevant results in your top retrieved chunks, leading to better answers when you pass these chunks to Claude for final response generation.\n",
    "\n",
    "The improvement is most pronounced at Pass@5, where precision matters most—suggesting that contextualized chunks aren't just retrieved more often, but rank higher when relevant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contextual BM25: Hybrid Search\n",
    "\n",
    "Contextual embeddings alone improved our Pass@10 from 87% to 92%. We can push performance even higher by combining semantic search with keyword-based search using **Contextual BM25**—a hybrid approach that reduces retrieval failure rates further.\n",
    "\n",
    "### Why Hybrid Search?\n",
    "\n",
    "Semantic search excels at understanding meaning and context, but can miss exact keyword matches. BM25 (a probabilistic keyword ranking algorithm) excels at finding specific terms, but lacks semantic understanding. By combining both, we get the best of both worlds:\n",
    "\n",
    "- **Semantic search**: Captures conceptual similarity and paraphrases\n",
    "- **BM25**: Catches exact terminology, function names, and specific phrases\n",
    "- **Reciprocal Rank Fusion**: Intelligently merges results from both sources\n",
    "\n",
    "### What is BM25?\n",
    "\n",
    "BM25 is a probabilistic ranking function that improves upon TF-IDF by accounting for document length and term saturation. It's widely used in production search engines (including Elasticsearch) for its effectiveness at ranking keyword relevance. For technical details, see [this blog post](https://www.elastic.co/blog/practical-bm25-part-2-the-bm25-algorithm-and-its-variables).\n",
    "\n",
    "Instead of only searching the raw chunk content, we search both the chunk *and* the contextual description we generated earlier. This means BM25 can match keywords in either the original text or the explanatory context.\n",
    "\n",
    "### Setup: Running Elasticsearch\n",
    "\n",
    "Before running the code below, you'll need Elasticsearch running locally. The easiest way is with Docker:\n",
    "\n",
    "```bash\n",
    "docker run -d --name elasticsearch -p 9200:9200 -p 9300:9300 \\\n",
    "  -e \"discovery.type=single-node\" \\\n",
    "  -e \"xpack.security.enabled=false\" \\\n",
    "  elasticsearch:9.2.0\n",
    "```\n",
    "\n",
    "## Troubleshooting:\n",
    "- Verify it's running: docker ps | grep elasticsearch\n",
    "- If port 9200 is in use: docker stop elasticsearch && docker rm elasticsearch\n",
    "- Check logs if issues occur: docker logs elasticsearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How the Hybrid Search Works\n",
    "\n",
    "The retrieve_advanced function below implements a three-step process:\n",
    "\n",
    "1. Retrieve candidates: Get top 150 results from both semantic search and BM25\n",
    "2. Score fusion: Combine rankings using weighted Reciprocal Rank Fusion\n",
    "    - Default: 80% weight to semantic search, 20% to BM25\n",
    "    - These weights are tunable—experiment to optimize for your use case\n",
    "3. Return top-k: Select the highest-scoring results after fusion\n",
    "\n",
    "The weighting system lets you balance between semantic understanding and keyword precision based on your data characteristics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from typing import Any\n",
    "\n",
    "from elasticsearch import Elasticsearch\n",
    "from elasticsearch.helpers import bulk\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "class ElasticsearchBM25:\n",
    "    def __init__(self, index_name: str = \"contextual_bm25_index\"):\n",
    "        self.es_client = Elasticsearch(\"http://localhost:9200\")\n",
    "        self.index_name = index_name\n",
    "        self.create_index()\n",
    "\n",
    "    def create_index(self):\n",
    "        index_settings = {\n",
    "            \"settings\": {\n",
    "                \"analysis\": {\"analyzer\": {\"default\": {\"type\": \"english\"}}},\n",
    "                \"similarity\": {\"default\": {\"type\": \"BM25\"}},\n",
    "                \"index.queries.cache.enabled\": False,\n",
    "            },\n",
    "            \"mappings\": {\n",
    "                \"properties\": {\n",
    "                    \"content\": {\"type\": \"text\", \"analyzer\": \"english\"},\n",
    "                    \"contextualized_content\": {\"type\": \"text\", \"analyzer\": \"english\"},\n",
    "                    \"doc_id\": {\"type\": \"keyword\", \"index\": False},\n",
    "                    \"chunk_id\": {\"type\": \"keyword\", \"index\": False},\n",
    "                    \"original_index\": {\"type\": \"integer\", \"index\": False},\n",
    "                }\n",
    "            },\n",
    "        }\n",
    "\n",
    "        # Change this line - remove 'body=' parameter\n",
    "        if not self.es_client.indices.exists(index=self.index_name):\n",
    "            self.es_client.indices.create(\n",
    "                index=self.index_name,\n",
    "                settings=index_settings[\"settings\"],\n",
    "                mappings=index_settings[\"mappings\"],\n",
    "            )\n",
    "            print(f\"Created index: {self.index_name}\")\n",
    "\n",
    "    def index_documents(self, documents: list[dict[str, Any]]):\n",
    "        actions = [\n",
    "            {\n",
    "                \"_index\": self.index_name,\n",
    "                \"_source\": {\n",
    "                    \"content\": doc[\"original_content\"],\n",
    "                    \"contextualized_content\": doc[\"contextualized_content\"],\n",
    "                    \"doc_id\": doc[\"doc_id\"],\n",
    "                    \"chunk_id\": doc[\"chunk_id\"],\n",
    "                    \"original_index\": doc[\"original_index\"],\n",
    "                },\n",
    "            }\n",
    "            for doc in documents\n",
    "        ]\n",
    "        success, _ = bulk(self.es_client, actions)\n",
    "        self.es_client.indices.refresh(index=self.index_name)\n",
    "        return success\n",
    "\n",
    "    def search(self, query: str, k: int = 20) -> list[dict[str, Any]]:\n",
    "        self.es_client.indices.refresh(index=self.index_name)\n",
    "\n",
    "        # Change this - remove 'body=' and pass query directly\n",
    "        response = self.es_client.search(\n",
    "            index=self.index_name,\n",
    "            query={\n",
    "                \"multi_match\": {\n",
    "                    \"query\": query,\n",
    "                    \"fields\": [\"content\", \"contextualized_content\"],\n",
    "                }\n",
    "            },\n",
    "            size=k,\n",
    "        )\n",
    "\n",
    "        return [\n",
    "            {\n",
    "                \"doc_id\": hit[\"_source\"][\"doc_id\"],\n",
    "                \"original_index\": hit[\"_source\"][\"original_index\"],\n",
    "                \"content\": hit[\"_source\"][\"content\"],\n",
    "                \"contextualized_content\": hit[\"_source\"][\"contextualized_content\"],\n",
    "                \"score\": hit[\"_score\"],\n",
    "            }\n",
    "            for hit in response[\"hits\"][\"hits\"]\n",
    "        ]\n",
    "\n",
    "\n",
    "def create_elasticsearch_bm25_index(db: ContextualVectorDB):\n",
    "    es_bm25 = ElasticsearchBM25()\n",
    "    es_bm25.index_documents(db.metadata)\n",
    "    return es_bm25\n",
    "\n",
    "\n",
    "def retrieve_advanced(\n",
    "    query: str,\n",
    "    db: ContextualVectorDB,\n",
    "    es_bm25: ElasticsearchBM25,\n",
    "    k: int,\n",
    "    semantic_weight: float = 0.8,\n",
    "    bm25_weight: float = 0.2,\n",
    "):\n",
    "    num_chunks_to_recall = 150\n",
    "\n",
    "    # Semantic search\n",
    "    semantic_results = db.search(query, k=num_chunks_to_recall)\n",
    "    ranked_chunk_ids = [\n",
    "        (result[\"metadata\"][\"doc_id\"], result[\"metadata\"][\"original_index\"])\n",
    "        for result in semantic_results\n",
    "    ]\n",
    "\n",
    "    # BM25 search using Elasticsearch\n",
    "    bm25_results = es_bm25.search(query, k=num_chunks_to_recall)\n",
    "    ranked_bm25_chunk_ids = [\n",
    "        (result[\"doc_id\"], result[\"original_index\"]) for result in bm25_results\n",
    "    ]\n",
    "\n",
    "    # Combine results\n",
    "    chunk_ids = list(set(ranked_chunk_ids + ranked_bm25_chunk_ids))\n",
    "    chunk_id_to_score = {}\n",
    "\n",
    "    # Initial scoring with weights\n",
    "    for chunk_id in chunk_ids:\n",
    "        score = 0\n",
    "        if chunk_id in ranked_chunk_ids:\n",
    "            index = ranked_chunk_ids.index(chunk_id)\n",
    "            score += semantic_weight * (1 / (index + 1))  # Weighted 1/n scoring for semantic\n",
    "        if chunk_id in ranked_bm25_chunk_ids:\n",
    "            index = ranked_bm25_chunk_ids.index(chunk_id)\n",
    "            score += bm25_weight * (1 / (index + 1))  # Weighted 1/n scoring for BM25\n",
    "        chunk_id_to_score[chunk_id] = score\n",
    "\n",
    "    # Sort chunk IDs by their scores in descending order\n",
    "    sorted_chunk_ids = sorted(\n",
    "        chunk_id_to_score.keys(), key=lambda x: (chunk_id_to_score[x], x[0], x[1]), reverse=True\n",
    "    )\n",
    "\n",
    "    # Assign new scores based on the sorted order\n",
    "    for index, chunk_id in enumerate(sorted_chunk_ids):\n",
    "        chunk_id_to_score[chunk_id] = 1 / (index + 1)\n",
    "\n",
    "    # Prepare the final results\n",
    "    final_results = []\n",
    "    semantic_count = 0\n",
    "    bm25_count = 0\n",
    "    for chunk_id in sorted_chunk_ids[:k]:\n",
    "        chunk_metadata = next(\n",
    "            chunk\n",
    "            for chunk in db.metadata\n",
    "            if chunk[\"doc_id\"] == chunk_id[0] and chunk[\"original_index\"] == chunk_id[1]\n",
    "        )\n",
    "        is_from_semantic = chunk_id in ranked_chunk_ids\n",
    "        is_from_bm25 = chunk_id in ranked_bm25_chunk_ids\n",
    "        final_results.append(\n",
    "            {\n",
    "                \"chunk\": chunk_metadata,\n",
    "                \"score\": chunk_id_to_score[chunk_id],\n",
    "                \"from_semantic\": is_from_semantic,\n",
    "                \"from_bm25\": is_from_bm25,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        if is_from_semantic and not is_from_bm25:\n",
    "            semantic_count += 1\n",
    "        elif is_from_bm25 and not is_from_semantic:\n",
    "            bm25_count += 1\n",
    "        else:  # it's in both\n",
    "            semantic_count += 0.5\n",
    "            bm25_count += 0.5\n",
    "\n",
    "    return final_results, semantic_count, bm25_count\n",
    "\n",
    "\n",
    "def evaluate_db_advanced(\n",
    "    db: ContextualVectorDB,\n",
    "    original_jsonl_path: str,\n",
    "    k_values: list[int] = None,\n",
    "    db_name: str = \"Hybrid Search\",\n",
    "):\n",
    "    \"\"\"\n",
    "    Evaluate hybrid search (semantic + BM25) at multiple k values with formatted results.\n",
    "\n",
    "    Args:\n",
    "        db: ContextualVectorDB instance\n",
    "        original_jsonl_path: Path to evaluation dataset\n",
    "        k_values: List of k values to evaluate (default: [5, 10, 20])\n",
    "        db_name: Name for the evaluation display\n",
    "\n",
    "    Returns:\n",
    "        Dict mapping k values to their results and source breakdowns\n",
    "    \"\"\"\n",
    "    if k_values is None:\n",
    "        k_values = [5, 10, 20]\n",
    "    original_data = load_jsonl(original_jsonl_path)\n",
    "    es_bm25 = create_elasticsearch_bm25_index(db)\n",
    "    results = {}\n",
    "\n",
    "    print(f\"{'=' * 70}\")\n",
    "    print(f\"Evaluation Results: {db_name}\")\n",
    "    print(f\"{'=' * 70}\\n\")\n",
    "\n",
    "    try:\n",
    "        # Warm-up queries\n",
    "        warm_up_queries = original_data[:10]\n",
    "        for query_item in warm_up_queries:\n",
    "            _ = retrieve_advanced(query_item[\"query\"], db, es_bm25, k_values[0])\n",
    "\n",
    "        for k in k_values:\n",
    "            print(f\"Evaluating Pass@{k}...\")\n",
    "\n",
    "            total_score = 0\n",
    "            total_semantic_count = 0\n",
    "            total_bm25_count = 0\n",
    "            total_results = 0\n",
    "\n",
    "            for query_item in tqdm(original_data, desc=f\"Pass@{k}\"):\n",
    "                query = query_item[\"query\"]\n",
    "                golden_chunk_uuids = query_item[\"golden_chunk_uuids\"]\n",
    "\n",
    "                golden_contents = []\n",
    "                for doc_uuid, chunk_index in golden_chunk_uuids:\n",
    "                    golden_doc = next(\n",
    "                        (doc for doc in query_item[\"golden_documents\"] if doc[\"uuid\"] == doc_uuid),\n",
    "                        None,\n",
    "                    )\n",
    "                    if golden_doc:\n",
    "                        golden_chunk = next(\n",
    "                            (\n",
    "                                chunk\n",
    "                                for chunk in golden_doc[\"chunks\"]\n",
    "                                if chunk[\"index\"] == chunk_index\n",
    "                            ),\n",
    "                            None,\n",
    "                        )\n",
    "                        if golden_chunk:\n",
    "                            golden_contents.append(golden_chunk[\"content\"].strip())\n",
    "\n",
    "                if not golden_contents:\n",
    "                    continue\n",
    "\n",
    "                retrieved_docs, semantic_count, bm25_count = retrieve_advanced(\n",
    "                    query, db, es_bm25, k\n",
    "                )\n",
    "\n",
    "                chunks_found = 0\n",
    "                for golden_content in golden_contents:\n",
    "                    for doc in retrieved_docs[:k]:\n",
    "                        retrieved_content = doc[\"chunk\"][\"original_content\"].strip()\n",
    "                        if retrieved_content == golden_content:\n",
    "                            chunks_found += 1\n",
    "                            break\n",
    "\n",
    "                query_score = chunks_found / len(golden_contents)\n",
    "                total_score += query_score\n",
    "\n",
    "                total_semantic_count += semantic_count\n",
    "                total_bm25_count += bm25_count\n",
    "                total_results += len(retrieved_docs)\n",
    "\n",
    "            total_queries = len(original_data)\n",
    "            average_score = total_score / total_queries\n",
    "            pass_at_n = average_score * 100\n",
    "\n",
    "            semantic_percentage = (\n",
    "                (total_semantic_count / total_results) * 100 if total_results > 0 else 0\n",
    "            )\n",
    "            bm25_percentage = (total_bm25_count / total_results) * 100 if total_results > 0 else 0\n",
    "\n",
    "            results[k] = {\n",
    "                \"pass_at_n\": pass_at_n,\n",
    "                \"average_score\": average_score,\n",
    "                \"total_queries\": total_queries,\n",
    "                \"semantic_percentage\": semantic_percentage,\n",
    "                \"bm25_percentage\": bm25_percentage,\n",
    "            }\n",
    "\n",
    "            print(f\"Pass@{k}: {pass_at_n:.2f}%\")\n",
    "            print(f\"Semantic: {semantic_percentage:.1f}% | BM25: {bm25_percentage:.1f}%\\n\")\n",
    "\n",
    "        # Print summary table\n",
    "        print(f\"{'=' * 70}\")\n",
    "        print(f\"{'Metric':<12} {'Pass Rate':<12} {'Score':<12} {'Semantic':<12} {'BM25':<12}\")\n",
    "        print(f\"{'-' * 70}\")\n",
    "        for k in k_values:\n",
    "            r = results[k]\n",
    "            print(\n",
    "                f\"{'Pass@' + str(k):<12} {r['pass_at_n']:>10.2f}% {r['average_score']:>10.4f} \"\n",
    "                f\"{r['semantic_percentage']:>10.1f}% {r['bm25_percentage']:>10.1f}%\"\n",
    "            )\n",
    "        print(f\"{'=' * 70}\\n\")\n",
    "\n",
    "        return results\n",
    "\n",
    "    finally:\n",
    "        # Delete the Elasticsearch index\n",
    "        if es_bm25.es_client.indices.exists(index=es_bm25.index_name):\n",
    "            es_bm25.es_client.indices.delete(index=es_bm25.index_name)\n",
    "            print(f\"Deleted Elasticsearch index: {es_bm25.index_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created index: contextual_bm25_index\n",
      "======================================================================\n",
      "Evaluation Results: Contextual BM25 Hybrid Search\n",
      "======================================================================\n",
      "\n",
      "Evaluating Pass@5...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@5: 100%|██████████| 248/248 [00:05<00:00, 41.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@5: 88.86%\n",
      "Semantic: 54.6% | BM25: 45.4%\n",
      "\n",
      "Evaluating Pass@10...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@10: 100%|██████████| 248/248 [00:05<00:00, 42.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@10: 92.31%\n",
      "Semantic: 57.6% | BM25: 42.4%\n",
      "\n",
      "Evaluating Pass@20...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@20: 100%|██████████| 248/248 [00:05<00:00, 42.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@20: 95.23%\n",
      "Semantic: 60.8% | BM25: 39.2%\n",
      "\n",
      "======================================================================\n",
      "Metric       Pass Rate    Score        Semantic     BM25        \n",
      "----------------------------------------------------------------------\n",
      "Pass@5            88.86%     0.8886       54.6%       45.4%\n",
      "Pass@10           92.31%     0.9231       57.6%       42.4%\n",
      "Pass@20           95.23%     0.9523       60.8%       39.2%\n",
      "======================================================================\n",
      "\n",
      "Deleted Elasticsearch index: contextual_bm25_index\n"
     ]
    }
   ],
   "source": [
    "results = evaluate_db_advanced(\n",
    "    contextual_db,\n",
    "    \"data/evaluation_set.jsonl\",\n",
    "    k_values=[5, 10, 20],\n",
    "    db_name=\"Contextual BM25 Hybrid Search\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reranking\n",
    "\n",
    "We've achieved strong results with hybrid search (93.21% Pass@10), but there's one more technique that can squeeze out additional performance: **reranking**.\n",
    "\n",
    "### What is Reranking?\n",
    "\n",
    "Reranking is a two-stage retrieval approach:\n",
    "\n",
    "1. **Stage 1 - Broad Retrieval**: Cast a wide net by retrieving more candidates than you need (e.g., retrieve 100 chunks)\n",
    "2. **Stage 2 - Precise Selection**: Use a specialized reranking model to score these candidates and select only the top-k most relevant ones\n",
    "\n",
    "**Why does this work?** Initial retrieval methods (embeddings, BM25) are optimized for speed across millions of documents. Reranking models are slower but more accurate—they can afford to do deeper analysis on a smaller candidate set. This creates a speed/accuracy trade-off that works well in practice.\n",
    "\n",
    "### Our Reranking Approach\n",
    "\n",
    "For this example, we'll use a simpler reranking pipeline that builds on contextual embeddings alone (not the full hybrid search). Here's the process:\n",
    "\n",
    "1. **Over-retrieve**: Get 10x more results than needed (e.g., retrieve 100 chunks when we need 10)\n",
    "2. **Rerank with Cohere**: Use Cohere's `rerank-english-v3.0` model to score all candidates\n",
    "3. **Select top-k**: Return only the highest-scoring results\n",
    "\n",
    "The reranking model has access to both the original chunk content and the contextual descriptions we generated, giving it rich information to make precise relevance judgments.\n",
    "\n",
    "### Expected Performance\n",
    "\n",
    "Adding reranking delivers a modest but meaningful improvement:\n",
    "- **Without reranking**: 92.34% Pass@10 (contextual embeddings alone)\n",
    "- **With reranking**: ~95% Pass@10 (additional 2-3% gain)\n",
    "\n",
    "This might seem small, but in production systems, reducing failures from 7.66% to ~5% can significantly improve user experience. The trade-off is query latency—reranking adds ~100-200ms per query depending on candidate set size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from collections.abc import Callable\n",
    "from typing import Any\n",
    "\n",
    "import cohere\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def evaluate_db_rerank(\n",
    "    db, original_jsonl_path: str, k_values: list[int] = None, db_name: str = \"Reranking\"\n",
    "):\n",
    "    \"\"\"\n",
    "    Evaluate reranking performance at multiple k values with formatted results.\n",
    "\n",
    "    Args:\n",
    "        db: ContextualVectorDB instance\n",
    "        original_jsonl_path: Path to evaluation dataset\n",
    "        k_values: List of k values to evaluate (default: [5, 10, 20])\n",
    "        db_name: Name for the evaluation display\n",
    "\n",
    "    Returns:\n",
    "        Dict mapping k values to their results\n",
    "    \"\"\"\n",
    "    if k_values is None:\n",
    "        k_values = [5, 10, 20]\n",
    "    original_data = load_jsonl(original_jsonl_path)\n",
    "    co = cohere.Client(os.getenv(\"COHERE_API_KEY\"))\n",
    "    results = {}\n",
    "\n",
    "    print(f\"{'=' * 60}\")\n",
    "    print(f\"Evaluation Results: {db_name}\")\n",
    "    print(f\"{'=' * 60}\\n\")\n",
    "\n",
    "    for k in k_values:\n",
    "        print(f\"Evaluating Pass@{k} with reranking...\")\n",
    "\n",
    "        total_score = 0\n",
    "        total_queries = len(original_data)\n",
    "\n",
    "        for query_item in tqdm(original_data, desc=f\"Pass@{k}\"):\n",
    "            query = query_item[\"query\"]\n",
    "            golden_chunk_uuids = query_item[\"golden_chunk_uuids\"]\n",
    "\n",
    "            # Find golden contents\n",
    "            golden_contents = []\n",
    "            for doc_uuid, chunk_index in golden_chunk_uuids:\n",
    "                golden_doc = next(\n",
    "                    (doc for doc in query_item[\"golden_documents\"] if doc[\"uuid\"] == doc_uuid), None\n",
    "                )\n",
    "                if golden_doc:\n",
    "                    golden_chunk = next(\n",
    "                        (chunk for chunk in golden_doc[\"chunks\"] if chunk[\"index\"] == chunk_index),\n",
    "                        None,\n",
    "                    )\n",
    "                    if golden_chunk:\n",
    "                        golden_contents.append(golden_chunk[\"content\"].strip())\n",
    "\n",
    "            if not golden_contents:\n",
    "                continue\n",
    "\n",
    "            # Retrieve and rerank\n",
    "            semantic_results = db.search(query, k=k * 10)\n",
    "\n",
    "            # Prepare documents for reranking\n",
    "            documents = [\n",
    "                f\"{res['metadata']['original_content']}\\n\\nContext: {res['metadata']['contextualized_content']}\"\n",
    "                for res in semantic_results\n",
    "            ]\n",
    "\n",
    "            # Rerank\n",
    "            rerank_response = co.rerank(\n",
    "                model=\"rerank-english-v3.0\", query=query, documents=documents, top_n=k\n",
    "            )\n",
    "            time.sleep(0.1)  # Rate limiting\n",
    "\n",
    "            # Get final results\n",
    "            retrieved_docs = []\n",
    "            for r in rerank_response.results:\n",
    "                original_result = semantic_results[r.index]\n",
    "                retrieved_docs.append(\n",
    "                    {\"chunk\": original_result[\"metadata\"], \"score\": r.relevance_score}\n",
    "                )\n",
    "\n",
    "            # Check if golden chunks are in results\n",
    "            chunks_found = 0\n",
    "            for golden_content in golden_contents:\n",
    "                for doc in retrieved_docs[:k]:\n",
    "                    retrieved_content = doc[\"chunk\"][\"original_content\"].strip()\n",
    "                    if retrieved_content == golden_content:\n",
    "                        chunks_found += 1\n",
    "                        break\n",
    "\n",
    "            query_score = chunks_found / len(golden_contents)\n",
    "            total_score += query_score\n",
    "\n",
    "        average_score = total_score / total_queries\n",
    "        pass_at_n = average_score * 100\n",
    "\n",
    "        results[k] = {\n",
    "            \"pass_at_n\": pass_at_n,\n",
    "            \"average_score\": average_score,\n",
    "            \"total_queries\": total_queries,\n",
    "        }\n",
    "\n",
    "        print(f\"Pass@{k}: {pass_at_n:.2f}%\")\n",
    "        print(f\"Average Score: {average_score:.4f}\\n\")\n",
    "\n",
    "    # Print summary table\n",
    "    print(f\"{'=' * 60}\")\n",
    "    print(f\"{'Metric':<15} {'Pass Rate':<15} {'Score':<15}\")\n",
    "    print(f\"{'-' * 60}\")\n",
    "    for k in k_values:\n",
    "        pass_rate = f\"{results[k]['pass_at_n']:.2f}%\"\n",
    "        score = f\"{results[k]['average_score']:.4f}\"\n",
    "        print(f\"{'Pass@' + str(k):<15} {pass_rate:<15} {score:<15}\")\n",
    "    print(f\"{'=' * 60}\\n\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================================================\n",
      "Evaluation Results: Contextual Embeddings + Reranking\n",
      "============================================================\n",
      "\n",
      "Evaluating Pass@5 with reranking...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@5: 100%|██████████| 248/248 [01:40<00:00,  2.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@5: 92.15%\n",
      "Average Score: 0.9215\n",
      "\n",
      "Evaluating Pass@10 with reranking...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@10: 100%|██████████| 248/248 [02:29<00:00,  1.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@10: 95.26%\n",
      "Average Score: 0.9526\n",
      "\n",
      "Evaluating Pass@20 with reranking...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pass@20: 100%|██████████| 248/248 [03:03<00:00,  1.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pass@20: 97.45%\n",
      "Average Score: 0.9745\n",
      "\n",
      "============================================================\n",
      "Metric          Pass Rate       Score          \n",
      "------------------------------------------------------------\n",
      "Pass@5          92.15%          0.9215         \n",
      "Pass@10         95.26%          0.9526         \n",
      "Pass@20         97.45%          0.9745         \n",
      "============================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results = evaluate_db_rerank(\n",
    "    contextual_db,\n",
    "    \"data/evaluation_set.jsonl\",\n",
    "    k_values=[5, 10, 20],\n",
    "    db_name=\"Contextual Embeddings + Reranking\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Reranking delivers our strongest results, nearly eliminating retrieval failures. Let's look at how each technique built upon the previous one to achieve this improvement.\n",
    "\n",
    "Starting from our baseline RAG system at 87% Pass@10, we've climbed to over 95% by systematically applying advanced retrieval techniques. Each method addresses a different weakness: contextual embeddings solve the \"isolated chunk\" problem, hybrid search catches keyword-specific queries that embeddings miss, and reranking applies more sophisticated relevance scoring to refine the final selection.\n",
    "\n",
    "| Approach | Pass@5 | Pass@10 | Pass@20 |\n",
    "|----------|--------|---------|---------|\n",
    "| **Baseline RAG** | 80.92% | 87.15% | 90.06% |\n",
    "| **+ Contextual Embeddings** | 88.12% | 92.34% | 94.29% |\n",
    "| **+ Hybrid Search (BM25)** | 86.43% | 93.21% | 94.99% |\n",
    "| **+ Reranking** | 92.15% | 95.26% | 97.45% |\n",
    "\n",
    "**Key Takeaways:**\n",
    "\n",
    "1. **Contextual embeddings provided the largest single improvement** (+5-7 percentage points), validating that adding document-level context to chunks significantly improves retrieval quality. This technique alone gets you 90% of the way to optimal performance.\n",
    "\n",
    "2. **Reranking achieves the highest absolute performance**, reaching 95.26% Pass@10—meaning the correct chunk appears in the top 10 results for 95% of queries. This represents a **47% reduction in retrieval failures** compared to baseline RAG (from 12.85% failure rate down to 4.74%).\n",
    "\n",
    "3. **Trade-offs matter**: Each technique adds complexity and cost:\n",
    "   - Contextual embeddings: One-time ingestion cost (~$3 for this dataset with prompt caching)\n",
    "   - Hybrid search: Requires Elasticsearch infrastructure and maintenance\n",
    "   - Reranking: Adds 100-200ms query latency and per-query API costs (~$0.002 per query)\n",
    "\n",
    "4. **Choose your approach** based on your requirements:\n",
    "   - **High-volume, cost-sensitive**: Contextual embeddings alone (92% Pass@10, no per-query costs)\n",
    "   - **Maximum accuracy, latency-tolerant**: Full reranking pipeline (95% Pass@10, best precision)\n",
    "   - **Balanced production system**: Hybrid search for strong performance without per-query costs (93% Pass@10)\n",
    "\n",
    "For most production RAG systems, **contextual embeddings provide the best performance-to-cost ratio**, delivering 92% Pass@10 with only one-time ingestion costs. Hybrid search and reranking are available when you need that extra 2-3 percentage points of precision and can afford the additional infrastructure or query costs.\n",
    "\n",
    "### Next Steps and Key Takeaways\n",
    "\n",
    "1) We demonstrated how to use Contextual Embeddings to improve retrieval performance, then delivered additional improvements with Contextual BM25 and reranking.\n",
    "\n",
    "2) This example used codebases, but these methods also apply to other data types such as internal company knowledge bases, financial & legal content, educational content, and much more. \n",
    "\n",
    "3) If you are an AWS user, you can get started with the Lambda function in `contextual-rag-lambda-function`, and if you're a GCP user you can spin up your own Cloud Run instance and follow a similar pattern!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anthropic-cookbook (3.12.12)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
